import asyncio
import datetime
from typing import Dict, List

from fastapi import WebSocket

from backend.report_type import BasicReport, DetailedReport
from backend.chat import ChatAgentWithMemory

from gpt_researcher.utils.enum import ReportType, Tone
from multi_agents.main import run_research_task
from gpt_researcher.actions import stream_output  # 导入流式输出方法
from backend.server.server_utils import CustomLogsHandler


class WebSocketManager:
    """WebSocket 管理器：负责管理所有 websocket 连接、消息队列和流式推送"""

    def __init__(self):
        """初始化 WebSocketManager 类。"""
        self.active_connections: List[WebSocket] = []  # 活跃 websocket 连接列表
        self.sender_tasks: Dict[WebSocket, asyncio.Task] = {}  # 每个 websocket 的异步发送任务
        self.message_queues: Dict[WebSocket, asyncio.Queue] = {}  # 每个 websocket 的消息队列
        self.chat_agent = None  # 聊天代理

    async def start_sender(self, websocket: WebSocket):
        """启动消息发送任务，不断从队列取消息并推送到前端。"""
        queue = self.message_queues.get(websocket)
        if not queue:
            return

        while True:
            try:
                message = await queue.get()
                if message is None:  # 关闭信号
                    break
                
                if websocket in self.active_connections:
                    if message == "ping":
                        await websocket.send_text("pong")  # 心跳检测
                    else:
                        await websocket.send_text(message)  # 发送普通消息
                else:
                    break
            except Exception as e:
                print(f"发送任务出错: {e}")
                break

    async def connect(self, websocket: WebSocket):
        """接收并注册新的 websocket 连接。"""
        try:
            await websocket.accept()
            self.active_connections.append(websocket)
            self.message_queues[websocket] = asyncio.Queue()
            self.sender_tasks[websocket] = asyncio.create_task(
                self.start_sender(websocket))  # 启动异步发送任务
        except Exception as e:
            print(f"连接 websocket 出错: {e}")
            if websocket in self.active_connections:
                await self.disconnect(websocket)

    async def disconnect(self, websocket: WebSocket):
        """断开并清理 websocket 连接及其资源。"""
        if websocket in self.active_connections:
            self.active_connections.remove(websocket)
            if websocket in self.sender_tasks:
                self.sender_tasks[websocket].cancel()
                await self.message_queues[websocket].put(None)  # 通知发送任务退出
                del self.sender_tasks[websocket]
            if websocket in self.message_queues:
                del self.message_queues[websocket]
            try:
                await websocket.close()
            except:
                pass  # 连接可能已关闭

    async def start_streaming(self, task, report_type, report_source, source_urls, document_urls, tone, websocket, headers=None, query_domains=[]):
        """启动研究任务，并将输出通过 websocket 实时推送。"""
        tone = Tone[tone]
        # 可自定义 JSON 配置文件路径
        config_path = "default"
        report = await run_agent(task, report_type, report_source, source_urls, document_urls, tone, websocket, headers=headers, query_domains=query_domains, config_path=config_path)
        # 每次生成新报告时，创建新的聊天代理
        self.chat_agent = ChatAgentWithMemory(report, config_path, headers)
        return report

    async def chat(self, message, websocket):
        """基于研究结果的智能聊天功能。"""
        if self.chat_agent:
            await self.chat_agent.chat(message, websocket)
        else:
            await websocket.send_json({"type": "chat", "content": "知识为空，请先运行研究任务获取知识"})

"""
multi_agents：多智能体协作，适合复杂任务，内容最丰富。
DetailedReport：详细分析，结构严谨，适合需要深入研究的场景。
BasicReport：简明扼要，适合快速获取主题概要。
"""

# 统一的 agent 运行入口，支持多种报告类型
async def run_agent(task, report_type, report_source, source_urls, document_urls, tone: Tone, websocket, stream_output=stream_output, headers=None, query_domains=[], config_path="", return_researcher=False):
    """运行 agent，根据报告类型选择不同的研究流程，并通过 websocket 实时推送日志。
    
    参数说明:
    task: 研究任务/查询内容，用户想要研究的主题
    report_type: 报告类型，可选值: "multi_agents"(多智能体), "DetailedReport"(详细报告), "BasicReport"(基础报告)
    report_source: 数据来源，如"web"(网络)、"documents"(文档)等
    source_urls: 指定的网址列表，作为研究参考来源
    document_urls: 文档地址列表，用于从特定文档获取信息
    tone: 输出报告的语气风格，如Formal(正式), Balanced(平衡), Casual(随意)等
    websocket: WebSocket连接对象，用于实时推送研究进度和结果
    stream_output: 流式输出方法，默认使用gpt_researcher中的stream_output
    headers: HTTP请求头信息，用于API调用时传递认证等信息
    query_domains: 限定的搜索域名列表，限制研究范围
    config_path: 配置文件路径，默认为"default"
    return_researcher: 是否返回研究者对象，用于后续操作
    """
    # 为本次研究任务创建日志处理器
    logs_handler = CustomLogsHandler(websocket, task)

    # 根据报告类型初始化不同的研究流程
    if report_type == "multi_agents":
        # 多智能体模式：启动多个协作智能体进行复杂研究
        report = await run_research_task(
            query=task, 
            websocket=logs_handler,  # 使用日志处理器包装 websocket
            stream_output=stream_output, 
            tone=tone, 
            headers=headers
        )
        report = report.get("report", "")  # 从返回结果中提取报告内容

    elif report_type == ReportType.DetailedReport.value:
        # 详细报告模式：生成全面、深入的分析报告
        researcher = DetailedReport(
            query=task,
            query_domains=query_domains,
            report_type=report_type,
            report_source=report_source,
            source_urls=source_urls,
            document_urls=document_urls,
            tone=tone,
            config_path=config_path,
            websocket=logs_handler,  # 使用日志处理器包装 websocket
            headers=headers
        )
        report = await researcher.run()  # 执行研究并获取报告
        
    else:
        # 基础报告模式（默认）：生成简洁、重点突出的基础报告
        researcher = BasicReport(
            query=task,
            query_domains=query_domains,
            report_type=report_type,
            report_source=report_source,
            source_urls=source_urls,
            document_urls=document_urls,
            tone=tone,
            config_path=config_path,
            websocket=logs_handler,  # 使用日志处理器包装 websocket
            headers=headers
        )
        report = await researcher.run()  # 执行研究并获取报告

    # 如果不是多智能体模式且需要返回研究者对象，则返回报告和研究者对象的元组
    if report_type != "multi_agents" and return_researcher:
        return report, researcher.gpt_researcher
    else:
        return report

    # 总结：run_agent 函数是研究任务的统一入口，根据用户选择的报告类型(多智能体/详细/基础)
    # 启动相应的研究流程，管理整个研究过程，并最终返回格式化的研究报告文本。
    # 该函数负责研究任务的分发、执行和结果整合，是系统核心功能的关键节点。
